// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ws

// Hub maintains the set of active clients and broadcasts messages to the
// clients.
type Hub struct {
	// Registered clients.
	clients map[*Client]bool

	// Registered clients by UserID.
	users map[string]*Client

	// Inbound messages from the clients.
	broadcast chan []byte

	// Register requests from the clients.
	register chan *Client

	// Unregister requests from clients.
	unregister chan *Client

	// Send message to a specific user.
	sendToUser chan struct {
		userID  string
		message []byte
	}
}

func NewHub() *Hub {
	return &Hub{
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		clients:    make(map[*Client]bool),
		users:      make(map[string]*Client),
		sendToUser: make(chan struct {
			userID  string
			message []byte
		}),
	}
}

func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.clients[client] = true
			if client.UserID != "" {
				h.users[client.UserID] = client
			}
		case client := <-h.unregister:
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				if client.UserID != "" {
					delete(h.users, client.UserID)
				}
				close(client.send)
			}
		case message := <-h.broadcast:
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clients, client)
					if client.UserID != "" {
						delete(h.users, client.UserID)
					}
				}
			}
		case msg := <-h.sendToUser:
			if client, ok := h.users[msg.userID]; ok {
				select {
				case client.send <- msg.message:
				default:
					close(client.send)
					delete(h.clients, client)
					delete(h.users, client.UserID)
				}
			}
		}
	}
}
